import torch
from torch import nn
from varflow.flows import FlowLayer
from varflow.utils import sum_except_batch


class Logits(FlowLayer):
    """A Variational sigmoid layer."""

    def __init__(self, encoder):
        super(Logits, self).__init__()
        self.encoder = encoder

    def log_prob(self, x):
        z, log_qz = self.encoder.sample_with_log_prob(context=x)
        log_px = -nn.functional.binary_cross_entropy_with_logits(z, x, reduction='none')
        log_px = sum_except_batch(log_px)
        return self.base_dist.log_prob(z) + log_px - log_qz

    def sample(self, num_samples):
        return self.base_dist.sample(num_samples).sigmoid()
